#include "LocalProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"

void LocalProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();


	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			GeometryGenerator geometryGenerator;
			GeometryGenerator::MeshData sphereData = geometryGenerator.CreateSphere(1.f, 10, 10);

			std::vector<glm::vec3> tempVertexs;
			tempVertexs.resize(sphereData.Vertices.size());
			for (uint32 i = 0; i < sphereData.Vertices.size(); ++i)
			{
				tempVertexs[i] = sphereData.Vertices[i].Position;
			}

			mVertexBuffer = VertexBuffer::Builder()
				.SetData(tempVertexs.data())
				.SetSize(tempVertexs.size() * sizeof(glm::vec3))
				.Build();

			mIndexBuffer = IndexBuffer::Builder()
				.SetData(sphereData.Indices32.data())
				.SetSize(sphereData.Indices32.size() * sizeof(uint32))
				.SetIndexCount(sphereData.Indices32.size())
				.Build();



			mDescriptor = RHIDescriptor::Builder()
				.PushShader<LocalVertexShader>()
				.PushShader<LocalFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();


			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();


			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build(this, "DepthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void LocalProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	std::vector<PushConstants> pushConstants;
	pushConstants.resize(mLocalScene->mBalls.size());
	for (int32 i = 0; i < mLocalScene->mBalls.size(); ++i)
	{
		pushConstants[i].model = glm::translate(glm::mat4(1.f), mLocalScene->mBalls[i]->position) * glm::scale(glm::mat4(1.f), glm::vec3(mLocalScene->mBalls[i]->radius));
		pushConstants[i].color = glm::vec4(mLocalScene->mBalls[i]->color, 1.f);
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo, pushConstants](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			std::vector<PushConstants> pushConstantsData = pushConstants;
			for (int32 i = 0; i < pushConstantsData.size(); ++i)
			{
				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);


				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(PushConstants), &pushConstantsData[i]);

				vkCmdDrawIndexed(commandBuffer, mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void LocalProcess::OnResize()
{
	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build(this, "DepthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
